"""
数据增强模块

提供超声心动图数据增强功能
"""

import cv2
import numpy as np
import torch
from typing import List, Dict, Tuple, Optional, Union, Any
import albumentations as A
from albumentations.pytorch import ToTensorV2
import random

from src.utils.logger import get_logger


class EchoAugmentation:
    """超声心动图数据增强类"""
    
    def __init__(self, 
                 image_size: Tuple[int, int] = (1024, 1024),
                 training: bool = True,
                 augmentation_config: Optional[Dict[str, Any]] = None):
        """
        初始化数据增强
        
        Args:
            image_size: 目标图像尺寸 (height, width)
            training: 是否为训练模式
            augmentation_config: 增强配置
        """
        self.image_size = image_size
        self.training = training
        self.logger = get_logger("EchoAugmentation")
        
        # 默认增强配置
        default_config = {
            'horizontal_flip': 0.5,
            'vertical_flip': 0.3,
            'rotation': 15,
            'scale': [0.8, 1.2],
            'brightness': 0.2,
            'contrast': 0.2,
            'noise': 0.1,
            'blur': 0.2,
            'elastic_transform': 0.3
        }
        
        if augmentation_config:
            default_config.update(augmentation_config)
        
        self.config = default_config
        self.transform = self._create_transform()
    
    def _create_transform(self) -> A.Compose:
        """创建增强变换"""
        transforms = []
        
        # 注意：Resize放在最后，确保所有变换后的尺寸都是有效的
        
        if self.training:
            # 几何变换
            if self.config.get('horizontal_flip', 0) > 0:
                transforms.append(A.HorizontalFlip(p=self.config['horizontal_flip']))
            
            if self.config.get('vertical_flip', 0) > 0:
                transforms.append(A.VerticalFlip(p=self.config['vertical_flip']))
            
            if self.config.get('rotation', 0) > 0:
                transforms.append(A.Rotate(
                    limit=self.config['rotation'],
                    p=0.5,
                    border_mode=cv2.BORDER_CONSTANT,
                    value=0
                ))
            
            # 暂时禁用RandomScale，因为它可能导致resize错误
            # if self.config.get('scale'):
            #     scale_limit = (self.config['scale'][1] - 1.0) * 100
            #     # 限制scale_limit以避免产生无效尺寸
            #     scale_limit = min(abs(scale_limit), 50)  # 最大50%
            #     transforms.append(A.RandomScale(
            #         scale_limit=scale_limit,
            #         p=0.5,
            #         interpolation=cv2.INTER_LINEAR
            #     ))
            
            # 颜色变换
            if self.config.get('brightness', 0) > 0 or self.config.get('contrast', 0) > 0:
                transforms.append(A.RandomBrightnessContrast(
                    brightness_limit=self.config['brightness'],
                    contrast_limit=self.config['contrast'],
                    p=0.5
                ))
            
            # 噪声和模糊
            if self.config.get('noise', 0) > 0:
                transforms.append(A.GaussNoise(
                    var_limit=(10.0, 50.0),
                    p=self.config['noise']
                ))
            
            if self.config.get('blur', 0) > 0:
                transforms.append(A.Blur(
                    blur_limit=3,
                    p=self.config['blur']
                ))
            
            # 弹性变换（模拟心脏跳动）
            if self.config.get('elastic_transform', 0) > 0:
                transforms.append(A.ElasticTransform(
                    alpha=1,
                    sigma=50,
                    alpha_affine=50,
                    p=self.config['elastic_transform'],
                    border_mode=cv2.BORDER_CONSTANT,
                    value=0
                ))

        # Resize放在最后，确保所有变换后的尺寸都是目标尺寸
        transforms.append(A.Resize(height=self.image_size[0], width=self.image_size[1], interpolation=cv2.INTER_LINEAR))

        # 归一化
        transforms.append(A.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ))

        # 转换为tensor
        transforms.append(ToTensorV2())

        return A.Compose(transforms, is_check_shapes=False)
    
    def __call__(self, image: np.ndarray, mask: Optional[np.ndarray] = None) -> Dict[str, torch.Tensor]:
        """
        应用数据增强
        
        Args:
            image: 输入图像 (H, W, C)
            mask: 输入掩码 (H, W) 或 (H, W, C)
            
        Returns:
            增强后的数据
        """
        if mask is not None:
            # 确保掩码是2D
            if len(mask.shape) == 3:
                if mask.shape[2] == 1:
                    mask = mask.squeeze(2)
                else:
                    # 如果是多通道mask，只取第一个通道
                    mask = mask[:, :, 0]
            
            # 确保image和mask尺寸匹配（使用tuple比较避免数组布尔值歧义）
            img_h, img_w = image.shape[:2]
            mask_h, mask_w = mask.shape[:2]
            
            # 验证尺寸有效性
            if img_h <= 0 or img_w <= 0 or mask_h <= 0 or mask_w <= 0:
                self.logger.warning(f"无效的尺寸: image {image.shape}, mask {mask.shape}")
                # 创建默认尺寸的mask
                mask = np.zeros((img_h if img_h > 0 else 1024, img_w if img_w > 0 else 1024), dtype=np.uint8)
            elif img_h != mask_h or img_w != mask_w:
                # 调整mask尺寸以匹配image
                try:
                    # 确保目标尺寸有效
                    if img_h > 0 and img_w > 0:
                        mask = cv2.resize(mask, (int(img_w), int(img_h)), interpolation=cv2.INTER_NEAREST)
                    else:
                        mask = np.zeros((img_h, img_w), dtype=np.uint8)
                except Exception as e:
                    self.logger.warning(f"调整mask尺寸失败: {e}, image shape: {image.shape}, mask shape: {mask.shape}")
                    # 如果调整失败，创建空mask
                    mask = np.zeros((img_h, img_w), dtype=np.uint8)
            
            # 确保mask数据类型正确
            if mask.dtype != np.uint8:
                mask = mask.astype(np.uint8)
            
            # 最终验证：确保图像和mask尺寸有效
            if img_h <= 0 or img_w <= 0:
                self.logger.warning(f"图像尺寸无效: {image.shape}，跳过数据增强")
                # 直接resize到目标尺寸
                image = cv2.resize(image, (self.image_size[1], self.image_size[0]), interpolation=cv2.INTER_LINEAR)
                mask = cv2.resize(mask, (self.image_size[1], self.image_size[0]), interpolation=cv2.INTER_NEAREST)
                # 手动归一化和转换
                image = image.astype(np.float32) / 255.0
                image = (image - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])
                image_tensor = torch.from_numpy(image).permute(2, 0, 1)
                mask_tensor = torch.from_numpy(mask).float()
                return {
                    'image': image_tensor,
                    'mask': mask_tensor
                }
            
            # 在应用变换前，确保图像和mask都是有效的numpy数组
            if not isinstance(image, np.ndarray) or image.size == 0:
                self.logger.error(f"无效的图像: {type(image)}, size: {image.size if hasattr(image, 'size') else 'unknown'}")
                image_tensor = torch.zeros((3, self.image_size[0], self.image_size[1]))
                mask_tensor = torch.zeros((self.image_size[0], self.image_size[1]))
                return {'image': image_tensor, 'mask': mask_tensor}
            
            if not isinstance(mask, np.ndarray) or mask.size == 0:
                self.logger.error(f"无效的mask: {type(mask)}, size: {mask.size if hasattr(mask, 'size') else 'unknown'}")
                image_tensor = torch.zeros((3, self.image_size[0], self.image_size[1]))
                mask_tensor = torch.zeros((self.image_size[0], self.image_size[1]))
                return {'image': image_tensor, 'mask': mask_tensor}
            
            try:
                transformed = self.transform(image=image, mask=mask)
            except Exception as e:
                # 如果数据增强失败，直接resize到目标尺寸并手动处理
                try:
                    self.logger.warning(f"数据增强失败，使用简单resize: {e}")
                    # 直接resize到目标尺寸
                    image_resized = cv2.resize(image, (self.image_size[1], self.image_size[0]), interpolation=cv2.INTER_LINEAR)
                    mask_resized = cv2.resize(mask, (self.image_size[1], self.image_size[0]), interpolation=cv2.INTER_NEAREST)
                    # 手动归一化
                    image_resized = image_resized.astype(np.float32) / 255.0
                    image_resized = (image_resized - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])
                    image_tensor = torch.from_numpy(image_resized).permute(2, 0, 1)
                    mask_tensor = torch.from_numpy(mask_resized).float()
                    return {
                        'image': image_tensor,
                        'mask': mask_tensor
                    }
                except Exception as e2:
                    self.logger.error(f"数据增强完全失败: {e2}, image shape: {image.shape}, mask shape: {mask.shape if mask is not None else None}")
                    # 返回空tensor
                    image_tensor = torch.zeros((3, self.image_size[0], self.image_size[1]))
                    mask_tensor = torch.zeros((self.image_size[0], self.image_size[1]))
                    return {
                        'image': image_tensor,
                        'mask': mask_tensor
                    }
            return {
                'image': transformed['image'],
                'mask': transformed['mask']
            }
        else:
            transformed = self.transform(image=image)
            return {
                'image': transformed['image']
            }


class TemporalAugmentation:
    """时序数据增强"""
    
    def __init__(self, 
                 frame_drop_prob: float = 0.1,
                 frame_repeat_prob: float = 0.1,
                 temporal_shift: bool = True):
        """
        初始化时序增强
        
        Args:
            frame_drop_prob: 帧丢弃概率
            frame_repeat_prob: 帧重复概率
            temporal_shift: 是否进行时序偏移
        """
        self.frame_drop_prob = frame_drop_prob
        self.frame_repeat_prob = frame_repeat_prob
        self.temporal_shift = temporal_shift
        self.logger = get_logger("TemporalAugmentation")
    
    def __call__(self, frames: List[np.ndarray], masks: Optional[List[np.ndarray]] = None) -> Tuple[List[np.ndarray], Optional[List[np.ndarray]]]:
        """
        应用时序增强
        
        Args:
            frames: 帧列表
            masks: 掩码列表（可选）
            
        Returns:
            增强后的帧和掩码
        """
        if len(frames) < 3:
            return frames, masks
        
        # 安全地处理masks（可能是list或numpy数组）
        has_masks = masks is not None
        if has_masks:
            try:
                if isinstance(masks, (list, tuple)):
                    has_masks = len(masks) > 0
                    augmented_masks = list(masks)  # 转换为列表以便操作
                else:
                    # numpy数组
                    has_masks = masks.size > 0
                    # 将numpy数组转换为列表
                    augmented_masks = [masks[i] for i in range(masks.shape[0])]
            except (TypeError, ValueError, AttributeError):
                has_masks = False
                augmented_masks = None
        else:
            augmented_masks = None
        
        augmented_frames = frames.copy()
        
        # 随机丢弃帧
        if random.random() < self.frame_drop_prob and len(augmented_frames) > 3:
            drop_idx = random.randint(0, len(augmented_frames) - 1)
            augmented_frames.pop(drop_idx)
            if has_masks and augmented_masks:
                augmented_masks.pop(drop_idx)
        
        # 随机重复帧
        if random.random() < self.frame_repeat_prob and len(augmented_frames) > 1:
            repeat_idx = random.randint(0, len(augmented_frames) - 1)
            augmented_frames.insert(repeat_idx, augmented_frames[repeat_idx])
            if has_masks and augmented_masks:
                augmented_masks.insert(repeat_idx, augmented_masks[repeat_idx])
        
        # 时序偏移
        if self.temporal_shift and len(augmented_frames) > 1:
            shift = random.randint(-len(augmented_frames) // 4, len(augmented_frames) // 4)
            if shift != 0:
                augmented_frames = augmented_frames[shift:] + augmented_frames[:shift]
                if has_masks and augmented_masks:
                    augmented_masks = augmented_masks[shift:] + augmented_masks[:shift]
        
        return augmented_frames, augmented_masks


class CardiacCycleAugmentation:
    """心脏周期增强"""
    
    def __init__(self, 
                 phase_shift_prob: float = 0.3,
                 cycle_length_variation: float = 0.1):
        """
        初始化心脏周期增强
        
        Args:
            phase_shift_prob: 相位偏移概率
            cycle_length_variation: 周期长度变化范围
        """
        self.phase_shift_prob = phase_shift_prob
        self.cycle_length_variation = cycle_length_variation
        self.logger = get_logger("CardiacCycleAugmentation")
    
    def detect_cardiac_phases(self, frames: List[np.ndarray]) -> List[float]:
        """
        检测心脏周期相位
        
        Args:
            frames: 帧列表
            
        Returns:
            相位列表 (0-1)
        """
        if len(frames) < 5:
            return [i / len(frames) for i in range(len(frames))]
        
        # 计算帧间差异
        diffs = []
        for i in range(1, len(frames)):
            if len(frames[i].shape) == 3:
                gray1 = cv2.cvtColor(frames[i-1], cv2.COLOR_BGR2GRAY)
                gray2 = cv2.cvtColor(frames[i], cv2.COLOR_BGR2GRAY)
            else:
                gray1, gray2 = frames[i-1], frames[i]
            
            diff = np.mean(np.abs(gray2.astype(float) - gray1.astype(float)))
            diffs.append(diff)
        
        # 寻找峰值（收缩期）
        from scipy.signal import find_peaks
        peaks, _ = find_peaks(diffs, height=np.mean(diffs), distance=3)
        
        if len(peaks) > 0:
            # 基于峰值计算相位
            cycle_length = len(frames) / len(peaks)
            phases = [(i % cycle_length) / cycle_length for i in range(len(frames))]
        else:
            # 线性相位
            phases = [i / len(frames) for i in range(len(frames))]
        
        return phases
    
    def __call__(self, frames: List[np.ndarray], masks: Optional[List[np.ndarray]] = None) -> Tuple[List[np.ndarray], Optional[List[np.ndarray]]]:
        """
        应用心脏周期增强
        
        Args:
            frames: 帧列表
            masks: 掩码列表（可选）
            
        Returns:
            增强后的帧和掩码
        """
        if len(frames) < 3:
            return frames, masks
        
        # 安全地处理masks（可能是list或numpy数组）
        has_masks = masks is not None
        if has_masks:
            try:
                if isinstance(masks, (list, tuple)):
                    has_masks = len(masks) > 0
                    masks_list = list(masks)  # 转换为列表以便操作
                else:
                    # numpy数组
                    has_masks = masks.size > 0
                    # 将numpy数组转换为列表
                    masks_list = [masks[i] for i in range(masks.shape[0])]
            except (TypeError, ValueError, AttributeError):
                has_masks = False
                masks_list = None
        else:
            masks_list = None
        
        # 检测心脏相位
        phases = self.detect_cardiac_phases(frames)
        
        # 相位偏移
        if random.random() < self.phase_shift_prob:
            shift_phase = random.uniform(0, 1)
            # 重新排列帧以模拟不同相位开始
            phase_diffs = [abs(p - shift_phase) for p in phases]
            sorted_indices = np.argsort(phase_diffs)
            
            frames = [frames[i] for i in sorted_indices]
            if has_masks and masks_list:
                masks_list = [masks_list[i] for i in sorted_indices]
        
        # 周期长度变化
        if random.random() < 0.5:
            variation = random.uniform(1 - self.cycle_length_variation, 
                                    1 + self.cycle_length_variation)
            new_length = int(len(frames) * variation)
            
            if new_length != len(frames):
                # 重采样帧
                indices = np.linspace(0, len(frames) - 1, new_length).astype(int)
                frames = [frames[i] for i in indices]
                if has_masks and masks_list:
                    masks_list = [masks_list[i] for i in indices]
        
        return frames, masks_list


def create_augmentation_pipeline(config: Dict[str, Any]) -> EchoAugmentation:
    """
    创建数据增强管道
    
    Args:
        config: 增强配置
        
    Returns:
        增强器
    """
    image_size = config.get('image_size', (1024, 1024))
    training = config.get('training', True)
    augmentation_config = config.get('augmentation', {})
    
    return EchoAugmentation(
        image_size=image_size,
        training=training,
        augmentation_config=augmentation_config
    )


def test_augmentation():
    """测试数据增强"""
    logger = get_logger("TestAugmentation")
    
    # 创建测试数据
    test_image = np.random.randint(0, 255, (512, 512, 3), dtype=np.uint8)
    test_mask = np.random.randint(0, 2, (512, 512), dtype=np.uint8)
    
    # 测试空间增强
    aug = EchoAugmentation(training=True)
    result = aug(test_image, test_mask)
    
    logger.info(f"原始图像尺寸: {test_image.shape}")
    logger.info(f"增强后图像尺寸: {result['image'].shape}")
    logger.info(f"增强后掩码尺寸: {result['mask'].shape}")
    
    # 测试时序增强
    frames = [test_image] * 10
    masks = [test_mask] * 10
    
    temporal_aug = TemporalAugmentation()
    aug_frames, aug_masks = temporal_aug(frames, masks)
    
    logger.info(f"原始帧数: {len(frames)}")
    logger.info(f"增强后帧数: {len(aug_frames)}")


if __name__ == "__main__":
    test_augmentation()
